import argparse
from trainer import Trainer
import torch
from dataset.utils import init_data
import logging

parse = argparse.ArgumentParser()
parse.add_argument('--epochs', type=int, help="训练批次", default=10000)
parse.add_argument('--batch_size', type=int, help="批次大小", default=99999999)
parse.add_argument('--season', type=str, choices=["春", "夏", "秋", "冬"], default="春")
parse.add_argument('--sheet_name', type=str, choices=["多云", "晴天", "少云", "阴天", "self"], default="self")
parse.add_argument('--varity', action="store_true", help="是否直接验证,默认否")
parse.add_argument('--step_back', action="store_true", help="是否进行逐步回归，默认否")
parse.add_argument('--varity_all', action='store_true', help="是否对bp和逐步回归同时验证，结果保存在一个图上")
parse.add_argument('--fontsize', type=int, default=16)
parse.add_argument('--extract_data', action='store_true', help="是否提取数据")
parse.add_argument('--in_features', type=int, help="输入维度", default=4)
parse.add_argument('--n', type=int, help="逐步回归选择维度个数", default=3)
parse.add_argument('--device', help="训练位置")
parse.add_argument('--logger', help="日志")
args = parse.parse_args()

args.device = torch.device("cuda:0")
logger = logging.getLogger("RegressionModel")
logger.setLevel(level=logging.DEBUG)
args.logger = logger

if __name__ == "__main__":

    if args.extract_data:
        init_data()

    trainer = Trainer(args)
    if args.varity:
        trainer.init_eval()
    seasons = ["春", "夏", "秋", "冬"]
    sheet_names = ["Sheet1"]
    for season in seasons:
        for sheet_name in sheet_names:

            args.season = season
            args.sheet_name = sheet_name
            trainer.reset_args(args)
            if args.varity or args.varity_all:
                trainer.varity()
            else:
                trainer.train()

    args.logger.info("运行成功!!!          \n")
